import mysql.connector
from mysql.connector import Error
import cv2  
import datetime

class ImageDB:
    def __init__(self, host, user, password, database):
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
        self.cursor = None

    def connect(self):
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database
            )
            self.cursor = self.connection.cursor()
            print("成功连接数据库")
        except Error as e:
            print(f"连接数据库失败: {e}")

    def insert_image(self, image_path, image_name):
        try:
            with open(image_path, 'rb') as file:
                image_data = file.read()

            sql = "INSERT INTO images (name, data) VALUES (%s, %s)"
            self.cursor.execute(sql, (image_name, image_data))
            self.connection.commit()
            print("图片已成功存入数据库")

        except Error as e:
            print(f"插入失败: {e}")
            self.connection.rollback()
        except FileNotFoundError:
            print(f"文件 {image_path} 不存在")


    def insert_detection_data(self, frame, classid, number):
        """插入检测结果到数据库
        参数:
            frame (numpy.ndarray): 图像数据
            classid (int): 类别ID
            number (int): 当前类别检测数量
        """
        try:
            if self.connection is None or not self.connection.is_connected():
                print("数据库未连接，正在尝试重新连接...")
                self.connect()

            # 1. 生成时间戳
            current_time = datetime.datetime.now()
            
            # 2. 转换numpy数组为二进制
            # 使用PNG格式压缩图像
            success, encoded_image = cv2.imencode('.png', frame)
            if not success:
                raise ValueError("无法编码图像数据")
            image_data = encoded_image.tobytes()

            # 3. 判断授粉条件
            ispollination = (classid == 0 and number > 3)

            # 4. 执行SQL插入
            sql = """INSERT INTO images 
                    (time, data, classid, ispollination)
                    VALUES (%s, %s, %s, %s)"""
            self.cursor.execute(sql, (
                current_time,
                image_data,
                classid,
                ispollination
            ))
            self.connection.commit()
            print("检测数据插入成功")

        except Error as e:
            print(f"数据库操作失败: {e}")
            if self.connection:
                self.connection.rollback()
        except Exception as e:
            print(f"数据处理异常: {e}")
            if self.connection:
                self.connection.rollback()

    def close(self):
        if self.connection:
            self.cursor.close()
            self.connection.close()
            print("数据库连接已关闭")

# 使用示例
if __name__ == "__main__":
    # 配置数据库连接参数
    db_config = {
        "host": "localhost",
        "user": "root",
        "password": "LiuJiaCheng666",
        "database": "image_db"
    }

    # 实例化类
    image_db = ImageDB(**db_config)
    image_db.connect()

    # 测试新接口（需要实际numpy图像数据）
    # 生成一个随机测试图像
    import numpy as np
    test_frame = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
    
    image_db.insert_detection_data(
        frame=test_frame,
        classid=0,
        number=5
    )

    image_db.close()